/*
* ====================================================================
* Copyright (c) 2013-2014 Neudesic.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that redistributions of source
* retain the above copyright notice, this list of conditions and the
* following disclaimer.
*
* This source file includes software written by Neudesic, LLC (http://www.neudesic.com).
*
* THIS SOFTWARE IS PROVIDED BY Neudesic, LLC ``AS IS'' AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY, NON-INFRINGEMENT, AND FITNESS
* FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  FURTHER, IN NO EVENT SHALL
* Neudesic, LLC (THE AUTHOR) OR ANY CONTRIBUTORS BE LIABLE FOR ANY DAMAGES
* WHATSOEVER THAT ARE RELATED TO THE SOFTWARE [INCLUDING BUT NOT LIMITED TO,
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* SUCH AS LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION, HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)] EVEN IF ADVISED OF THE POSSIBILITY
* OF SUCH DAMAGES.
*
* ====================================================================
*/
#include "DatagramSocketContainer.h"
#include "winrtutils.h"
#include "SocketContainer.h"
#include "ISocketContainer.h"
#include <winstring.h>
#include <vector>

using ABI::Windows::Foundation::AsyncStatus;
using ABI::Windows::Foundation::IAsyncAction;
using ABI::Windows::Foundation::IAsyncOperation;
using ABI::Windows::Foundation::IAsyncOperationCompletedHandler;
using ABI::Windows::Foundation::IAsyncOperationWithProgressCompletedHandler;
using ABI::Windows::Foundation::IAsyncActionCompletedHandler;
using ABI::Windows::Foundation::ITypedEventHandler;
using ABI::Windows::Networking::Sockets::DatagramSocket;
using ABI::Windows::Networking::Sockets::DatagramSocketMessageReceivedEventArgs;
using ABI::Windows::Networking::Sockets::IDatagramSocketControl;
using ABI::Windows::Networking::Sockets::IDatagramSocketInformation;
using ABI::Windows::Networking::Sockets::IDatagramSocketMessageReceivedEventArgs;
using ABI::Windows::Networking::IEndpointPair;
using ABI::Windows::Networking::IEndpointPairFactory;
using ABI::Windows::Storage::Streams::IBuffer;
using ABI::Windows::Storage::Streams::IInputStream;
using ABI::Windows::Storage::Streams::IDataWriter;
using Microsoft::WRL::Callback;
using Microsoft::WRL::Wrappers::HString;
using Microsoft::WRL::Wrappers::HStringReference;

namespace Neudesic {
	namespace Sockets {

		DatagramSocketContainer::DatagramSocketContainer()
			: SocketContainer(RuntimeClass_Windows_Networking_Sockets_DatagramSocket, SocketType::DatagramSocketType)
		{
			//::ShowInfoMessage(L"In DatagramSocketContainer::c'tor for container [%d].", this);
			messageReceivedToken.value = 0;
		}

		DatagramSocketContainer::~DatagramSocketContainer(void)
		{
			//::ShowInfoMessage(L"DatagramSocketContainer::d'tor for container [%d] socket [%d].", this, this->GetSocket());
		}

		HRESULT DatagramSocketContainer::Close()
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;
			SocketContainer::Close();
			if (this->messageReceivedToken.value > 0)
			{
				HRESULT hr = this->GetSocket()->remove_MessageReceived(this->messageReceivedToken);
				if (FAILED(hr))
				{
					//::ShowErrorMessageWithResult(hr, L"Failed removing MessageReceivedToken [%d] for container [%d] socket [%d].",
					//	this->messageReceivedToken.value, this, this->GetSocket());
				}
				else
				{
					::ShowInfoMessage(L"Successfully removed MessageReceivedToken [%d] for container [%d] socket [%d].",
						this->messageReceivedToken.value, this, this->GetSocket());
				}
			}
			this->messageReceivedToken.value = 0;
			::WSASetLastError(0);
			return S_OK;
		}

		HRESULT DatagramSocketContainer::BindSocket(const struct sockaddr* name)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;
			//The intent is to reserve or name the socket and listen for incoming data through the MessageReceived
			//event. The name parameter is a local address and port on which to listen. Unless the IDatagramSocket
			//is supposed to listen to a specific remote endpoint, ConnectAsync with a IDatagramSocket is not
			//necessary. If not connected, it will receive data from any remote endpoint.
			std::wstring host;
			std::wstring port;
			bool anyPort = this->GetPortAndHost(name, host, port);
			//::ShowInfoMessage(L"From SocketContainer::Bind:Trying to bind to IP address [%s] port [%s].", host.c_str(), port.c_str());
			ComPtr<IHostName> hostName;
			ComPtr<IHostNameFactory> hostNameFactory;
			// Create the hostname
			HRESULT hr = Windows::Foundation::GetActivationFactory(HStringReference(RuntimeClass_Windows_Networking_HostName).Get(), &hostNameFactory);
			if (SUCCEEDED(hr))
			{
				hr = hostNameFactory->CreateHostName(HStringReference(host.c_str()).Get(), &hostName);
				if (SUCCEEDED(hr))
				{
					while (true)
					{
						if (anyPort)
						{
							port = SocketContainer::GetDynamicPort();
						}
						::ShowInfoMessage(L"Trying to bind to [%s] port [%s] for container [%d] socket [%d].",
							host.c_str(), port.c_str(), this, this->GetSocket());
						hr = this->Bind(hostName, port.c_str());
						if (SUCCEEDED(hr) || !anyPort)	//If binding to a specific port, forwward the failure regardless
							break;
						//Getting here means that we are using dynamic ports. If the port is already in use, get another socket.
						if (SocketErrorStatus::SocketErrorStatus_AddressAlreadyInUse != ::GetSocketErrorStatus(hr))
							break;
					}
					SocketContainer::GetHostAndPortInformation(hostName, host, port, &this->bindHostAndPortInfo);
				}
			}
			return hr;
		}

		HRESULT DatagramSocketContainer::Bind(const ComPtr<IHostName>& hostName, const wchar_t* serviceName)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;

			//Before we bind, the container must do some pre work. Datagram sockets must subscribe to
			//the MessageReceived event. This must happen before we bind. trying to subscribe to the
			//events after the bind will fail with E_ILLEGAL_METHOD_CALL(0x8000000e). This basically
			//means that a method was called unexpectedly.
			HRESULT hr = PrepareForMessages();
			ComPtr<IDatagramSocket> listeningSocket = this->GetSocket();
			if (nullptr == listeningSocket)
			{
				::ShowWarningMessage(L"PrepareForBind returned a NULL listener for container [%d] socket [%d].", this, this->GetSocket());
			}
			else
			{
				AutoResetEvent bindCompletedEvent;
				ComPtr<IAsyncActionCompletedHandler> bindCompletedHandler = ::Callback<IAsyncActionCompletedHandler>(
					[this, &bindCompletedEvent](IAsyncAction *asyncAction, AsyncStatus asyncStatus)
				{
					HRESULT bindCompletedResult = asyncAction->GetResults();
					switch (asyncStatus)
					{
					case AsyncStatus::Canceled:
					case AsyncStatus::Error:
						if (!FAILED(bindCompletedResult))
							bindCompletedResult = E_FAIL;
						break;
					case AsyncStatus::Completed:
						this->SetContainerState(SocketContainerState::Bound);
						break;
					}
					::ContinueIfSucceeded(bindCompletedResult, L"Bind failure for container [%d] socket [%d]. AsyncStatus [%d].",
						this, this->GetSocket(), asyncStatus);
					bindCompletedEvent.Set();
					return bindCompletedResult;
				});
				ComPtr<IAsyncAction> asyncAction;
				//If a hostname is provided, use it. Otherwise just bind to the port.
				HostAndPortInformation bindHostAndPortInfo = this->bindHostAndPortInfo;
				hr = (bindHostAndPortInfo.host.length() <= 1)
					? listeningSocket->BindServiceNameAsync(HStringReference(serviceName).Get(), &asyncAction)
					: listeningSocket->BindEndpointAsync(hostName.Get(), HStringReference(serviceName).Get(), &asyncAction);
				if (::ContinueIfSucceeded(hr, L"Bind failed to [%s] port [%s].", bindHostAndPortInfo.host.c_str(), bindHostAndPortInfo.port.c_str()))
				{
					hr = asyncAction->put_Completed(bindCompletedHandler.Get());
					if (::ContinueIfSucceeded(hr, L"put_Completed failed."))
					{
						int result = this->WaitForEvent(bindCompletedEvent.Get(), TimeoutDuration::Bind, L"BindCompleted");
						if (WAIT_TIMEOUT == result)
							::ShowErrorMessage(L"Timeout waiting for bind to stream listener.");
						hr = asyncAction->GetResults();
						::ContinueIfSucceeded(hr, L"Bind failed for container [%d] socket [%d].", this, this->GetSocket());
					}
				}
				bindCompletedEvent.Close();
			}
			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::ConnectSocket(const ComPtr<IHostName>& hostName, const wchar_t* serviceName)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;

			ComPtr<IAsyncAction> connectAsyncAction;
			HRESULT hr = this->GetSocket()->ConnectAsync(hostName.Get(), HStringReference(serviceName).Get(), &connectAsyncAction);
			if (::ContinueIfSucceeded(hr, L"Could not connect to remote service name for container [%d] socket [%d].", this, this->GetSocket()))
			{
				this->SetConnectAction(connectAsyncAction);
				//Datagram sockets merely record the peers address. The connection is made when we send/receive. A
				//connect completion handler will never fire so not needed.
				this->SetConnectComplete();
			}
			return hr;
		}

		HRESULT DatagramSocketContainer::EventSelect(WSAEVENT hEventObject, long lNetworkEvents)
		{
			//A DatagramSocket will never raise FD_CONNECT. For SIP we start by doing a REGISTER. Let's
			//set the write bit since we know that the stack will be sending a REGISTER request.
			HRESULT hr = SocketContainer::EventSelect(hEventObject, lNetworkEvents);
			if (lNetworkEvents & FD_WRITE)
				this->SetNetworkEvent(FD_WRITE, FD_WRITE_BIT, hr);
			return hr;
		}

		HRESULT DatagramSocketContainer::GetSocketName(struct sockaddr* name, int* namelen) const
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;

			HRESULT hr = SOCKET_ERROR;
			//Like winsock, bind or connect must have been called. They will have saved the IP and port info.
			if (SocketContainerState::Bound == this->GetContainerState() || SocketContainerState::Connected == this->GetContainerState())
			{
				SocketInformation<IDatagramSocket, IDatagramSocketInformation> si(this->GetSocket());
				HostAndPortInformation info;
				si.GetLocalHostAndPortInfo(&info);
				//Is this an ANY bind address? If so, clear it. A bind should have already occured to a specific address.
				if (info.host.compare(L"::") == 0 || info.host.compare(L"0.0.0.0") == 0)
					info.host.clear();

				if (!info.host.length())
					info = this->bindHostAndPortInfo;
				//::ShowInfoMessage(L"GetSocketName::LocalCanonicalName [%s] LocalPort [%s] HostNameType [%d] for container [%d] socket [%d].",
				//	info.host.c_str(), info.port.c_str(), info.type, this, this->GetSocket());
				hr = this->SetSockAddrInCommon((void*)name, namelen, info.host, info.type, info.port);
			}
			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::Send(const char* buf, unsigned int len, const ComPtr<IHostName>& hostName, const wchar_t* port, UINT32* numberOfBytesSent)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;

			if (numberOfBytesSent)
				*numberOfBytesSent = 0;
			//::ShowInfoMessage(L"Send:Number of bytes to send [%d] for container [%d] socket [%d].", len, this, this->GetSocket());
			if (!len)
				return S_OK;

			HANDLE hWriteOutputStreamComplete = ::CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);
			auto writeOutputStreamCompleteHandler = ::Callback<IAsyncOperationWithProgressCompletedHandler<UINT32, UINT32>>(
				[this, numberOfBytesSent, hWriteOutputStreamComplete](IAsyncOperationWithProgress<UINT32, UINT32> *operation, AsyncStatus asyncStatus)
			{
				HRESULT writeOutputStreamCompleteResult = operation->GetResults(numberOfBytesSent);
				switch (asyncStatus)
				{
				case AsyncStatus::Canceled:
				case AsyncStatus::Error:
					if (!FAILED(writeOutputStreamCompleteResult))
						writeOutputStreamCompleteResult = E_FAIL;
					break;
				case AsyncStatus::Completed:
					writeOutputStreamCompleteResult = operation->GetResults(numberOfBytesSent);
					//::ShowInfoMessage(L"SendComplete:Number of bytes sent [%d] for container [%d] socket [%d]",
					//	*numberOfBytesSent, this, this->GetSocket());
					::ContinueIfSucceeded(writeOutputStreamCompleteResult, L"Get number of bytes sent failed for container [%d] socket [%d].", this, this->GetSocket());
					break;
				}
				::ContinueIfSucceeded(writeOutputStreamCompleteResult, L"WriteBytes failure for container [%d] socket [%d]. AsyncStatus [%d].",
					this, this->GetSocket(), asyncStatus);
				::SetEvent(hWriteOutputStreamComplete);
				return writeOutputStreamCompleteResult;
			});
#if !USE_ASYNC_IO
			HANDLE hSendComplete = ::CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);
#endif
			ComPtr<IDataWriter> writer;
			HRESULT hr = Windows::Foundation::ActivateInstance(HStringReference(RuntimeClass_Windows_Storage_Streams_DataWriter).Get(), &writer);
			if (::ContinueIfSucceeded(hr, L"Activate instance of data writer failed for container [%d] socket [%d].", this, this->GetSocket()))
			{
				hr = writer->WriteBytes(len, (BYTE*)buf);
				if (::ContinueIfSucceeded(hr, L"Write bytes failed for container [%d] socket [%d].", this, this->GetSocket()))
				{
					ComPtr<IBuffer> sendBuffer;
					hr = writer->DetachBuffer(&sendBuffer);
					if (::ContinueIfSucceeded(hr, L"Detach send buffer failed for container [%d] socket [%d].", this, this->GetSocket()))
					{
#if USE_ASYNC_IO
						auto outputStreamCompletedHandler = ::Callback<IAsyncOperationCompletedHandler<IOutputStream*>>(
							[this, sendBuffer, writeOutputStreamCompleteHandler, hWriteOutputStreamComplete]
#else
						auto outputStreamCompletedHandler = ::Callback<IAsyncOperationCompletedHandler<IOutputStream*>>(
							[this, sendBuffer, writeOutputStreamCompleteHandler, hWriteOutputStreamComplete, &hSendComplete]
#endif
						(IAsyncOperation<IOutputStream*> *operation, AsyncStatus asyncStatus)
						{
							ComPtr<IOutputStream> outputStream;
							HRESULT completedResult = operation->GetResults(&outputStream);
							ComPtr<IAsyncOperationWithProgress<UINT32, UINT32>> action;
							switch (asyncStatus)
							{
							case AsyncStatus::Canceled:
							case AsyncStatus::Error:
								if (!FAILED(completedResult))
									completedResult = E_FAIL;
								break;
							case AsyncStatus::Completed:
								if (::ContinueIfSucceeded(completedResult, L"Get output stream failed for container [%d] socket [%d].", this, this->GetSocket()))
								{
									//this->SetReceiveAsyncOperation(action);
									completedResult = outputStream->WriteAsync(sendBuffer.Get(), &action);
									if (::ContinueIfSucceeded(completedResult, L"Write async on output stream failed for container [%d] socket [%d].", this, this->GetSocket()))
									{
										completedResult = action->put_Completed(writeOutputStreamCompleteHandler.Get());
										if (::ContinueIfSucceeded(completedResult, L"Call to put_Completed failed for container [%d] socket [%d].", this, this->GetSocket()))
										{
											int result = this->WaitForEvent(hWriteOutputStreamComplete, TimeoutDuration::Send, L"WriteComplete");
											if (result == WAIT_TIMEOUT)
												::ShowErrorMessage(L"Timeout waiting for completion on WriteAsync.");
										}
									}
								}
								break;
							}
							::ContinueIfSucceeded(completedResult, L"OutputStream failure for container [%d] socket [%d]. AsyncStatus [%d].",
								this, this->GetSocket(), asyncStatus);
							::CloseHandle(hWriteOutputStreamComplete);
#if USE_ASYNC_IO
							this->SignalSendCompletion();
#else
							::SetEvent(hSendComplete);
#endif
							return completedResult;
						});
						ComPtr<IEndpointPairFactory> endpointFactory;
						ComPtr<IEndpointPair> endpointPair;

						hr = Windows::Foundation::GetActivationFactory(HStringReference(RuntimeClass_Windows_Networking_EndpointPair).Get(), &endpointFactory);
						if (::ContinueIfSucceeded(hr, L"Activate endpoint pair factory failed for container [%d] socket [%d].", this, this->GetSocket()))
						{
							hr = endpointFactory->CreateEndpointPair(nullptr, HStringReference(L"").Get(), hostName.Get(), HStringReference(port).Get(), &endpointPair);
							if (::ContinueIfSucceeded(hr, L"Create endpoint pair failed for container [%d] socket [%d].", this, this->GetSocket()))
							{
								ComPtr<IAsyncOperation<IOutputStream*>> operation;
								ComPtr<IDatagramSocket> socket = this->GetSocket();

								hr = socket->GetOutputStreamWithEndpointPairAsync(endpointPair.Get(), &operation);
								if (::ContinueIfSucceeded(hr, L"Get output stream async failed for container [%d] socket [%d].", this, this->GetSocket()))
								{
									hr = operation->put_Completed(outputStreamCompletedHandler.Get());
									if (::ContinueIfSucceeded(hr, L"Set OutputStream completion handler failed for socket [%d].", this))
									{
#if !USE_ASYNC_IO
										int result = this->WaitForEvent(hSendComplete, TimeoutDuration::Send, L"SendComplete");
										if (WAIT_TIMEOUT == result)
										{
											::ShowErrorMessage(L"Timeout waiting for SendCompleted.");
											hr = HRESULT_FROM_WIN32(result);
										}
#endif
									}
								}
							}
						}
					}
				}
			}

#if !USE_ASYNC_IO
			::CloseHandle(hSendComplete);
#endif

			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::SendTo(const char* buf, unsigned int len, int flags, const sockaddr* to, int tolen, UINT32* numberOfBytesSent)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;
			if (numberOfBytesSent)
				*numberOfBytesSent = 0;

			//::ShowInfoMessage(L"In DatagramSocketContainer::SendTo for socket [%d].", this);
			std::wstring host;
			std::wstring port;
			this->GetPortAndHost(to, host, port);

			ComPtr<IHostName> hostName;
			ComPtr<IHostNameFactory> hostNameFactory;

			// Create the hostname
			HRESULT hr = ABI::Windows::Foundation::GetActivationFactory(HStringReference(RuntimeClass_Windows_Networking_HostName).Get(), &hostNameFactory);
			if (::ContinueIfSucceeded(hr, L"Getting activation factory for hostname failed for socket [%d].", this))
			{
				hr = hostNameFactory->CreateHostName(HStringReference(host.c_str()).Get(), &hostName);
				if (::ContinueIfSucceeded(hr, L"Creating hostname failed for socket [%d].", this))
				{
					// Hand it off to the socket or do the connect and send
					hr = this->Send(buf, len, hostName, port.c_str(), numberOfBytesSent);
				}
			}
			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::ReceiveFrom(char *buf, int len, int flags, struct sockaddr *from, int *fromlen, LPDWORD lpNumberOfBytesRecvd)
		{
			if (!this->IsRunning())
				return INVALID_SOCKET;
			struct sockaddr_storage remoteAddress;
			int namelen = sizeof(sockaddr_storage);
			HRESULT hr = this->GetPeerName((sockaddr*)&remoteAddress, &namelen);
			if (::ContinueIfSucceeded(hr, L"Invalid remote address or port for socket [%d].", this))
			{
				*fromlen = min(namelen, *fromlen);;
				::memcpy_s(from, *fromlen, &remoteAddress, *fromlen);
			}
			//Do read last so that event is raised last.
			hr = this->InternalReceive(buf, len, flags, lpNumberOfBytesRecvd);	//Common receive
			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::SetSocketOptions(int level, int optname, const char* optval, int optlen)
		{
			//At this time there are no options used by Doubango that requires setting control options.

			HRESULT hr = S_OK;
			//HRESULT hr = SOCKET_ERROR;
			//if (this->IsRunning())
			//{
			//	//ComPtr<IDatagramSocketControl> control;
			//	//hr = this->GetSocket()->get_Control(&control);
			//}
			::WSASetLastError(FAILED(hr) ? WSAEFAULT : 0);
			return hr;
		}

		HRESULT DatagramSocketContainer::PrepareForMessages()
		{
			//::ShowInfoMessage(L"In PrepareForMessages for container [%d] socket [%d].", this, this->GetSocket());
			auto messageReceivedHandler = ::Callback<ITypedEventHandler<DatagramSocket*, DatagramSocketMessageReceivedEventArgs*>>(
				[this](IDatagramSocket *socket, IDatagramSocketMessageReceivedEventArgs *eventArguments)
			{
				if (!this->IsRunning())
					return (HRESULT)INVALID_SOCKET;

				//::ShowInfoMessage(L"Received message for container [%d] socket [%d].", this, socket);
				HRESULT hr = S_OK;
				ComPtr<IDataReader> reader;
				hr = eventArguments->GetDataReader(&reader);
				if (::ContinueIfSucceeded(hr, L"Could not get DataReader for container [%d] socket [%d].", this, this->GetSocket()))
				{
					this->PrepareDataReader(reader);
					ComPtr<IHostName> remoteHostName;
					hr = eventArguments->get_RemoteAddress(&remoteHostName);
					if (::ContinueIfSucceeded(hr, L"Could not get remote host name for container [%d] socket [%d].", this, socket))
					{
						HSTRING hport;
						hr = eventArguments->get_RemotePort(&hport);
						if (::ContinueIfSucceeded(hr, L"Could not get remote port name for container [%d] socket [%d].", this, socket))
						{
							::WindowsDeleteString(hport);
						}
					}
				}
#if USE_ASYNC_IO
				this->SignalReceiveComplete();
#else
				this->SetNetworkEvent(FD_READ, FD_READ_BIT, hr);
#endif
				return hr;
			});
			HRESULT hr = this->GetSocket()->add_MessageReceived(messageReceivedHandler.Get(), &this->messageReceivedToken);
			if (::ContinueIfSucceeded(hr, L"Failed to subscribe to message received for container [%d] socket [%d].", this, this->GetSocket()))
			{
				//::ShowInfoMessage(L"Added MessageReceivedToken [%d] for container [%d] socket [%d].", this->messageReceivedToken.value, this, this->GetSocket());
			}
			return hr;
		}
	}
}